import torch
import torch.nn as nn

test_input = torch.randint(0, 10, (4, 8))

"""
Dropout的作用机制
不是简单地将0.1的数值设置为0，而是以0.1的概率随机将输入张量中的部分元素置为0
dropout参数（这里是0.1）表示每个神经元被置零的概率
剩余元素会进行缩放，通常乘以1/(1-dropout)来保持整体期望值不变
---
在训练过程中，随机将10%的神经元输出置为0
其余90%的输出会乘以10/9（约1.11）进行缩放
在评估/测试模式下，dropout层不会产生任何效果（即不置零也不缩放）
---
防止过拟合：通过随机丢弃神经元，防止网络对特定神经元产生过度依赖
提高泛化能力：使模型在训练时更加鲁棒，提高在未见数据上的表现
因此，dropout是一种正则化技术，而不是简单地将特定数值替换为0
---
Dropout 层总是对其接收到的输入张量执行操作（也就是前一层的输出）
在训练模式下，它会随机将输入张量中的一些元素置为0
在推理模式下（eval mode），dropout 不会对数据做任何修改
"""


class PositionwiseFeedForward(nn.Module):
    """
    Transformer中的Feed Forward模块
    """

    def __init__(self, d_model=512, d_ff=2048, dropout=0.1):
        super(PositionwiseFeedForward, self).__init__()
        self.linear1 = nn.Linear(d_model, d_ff)
        self.activation = nn.ReLU()
        self.linear2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        return self.linear2(self.dropout(self.activation(self.linear1(x))))


class FeedForwardNet(nn.Module):
    def __init__(self):
        super(FeedForwardNet, self).__init__()
        self.embedding = nn.Embedding(10, 512)
        self.position_encoding = PositionwiseFeedForward(512, 2048, 0.1)
        self.linear = nn.Linear(512, 10)

    def forward(self, x):
        x = self.embedding(x)
        x = self.position_encoding(x)
        x = self.linear(x)
        return x


if __name__ == '__main__':
    model = FeedForwardNet()
    print(test_input)
    print(model(test_input))
